from pyspark.sql.session import SparkSession
from pyspark.sql.functions import *
from pyspark.ml.linalg import Vectors, VectorUDT

from pyspark.ml.classification import LogisticRegressionModel
from pyspark.sql.types import *

# 创建环境
spark = SparkSession.builder.getOrCreate()

# 1、读取图片数据
image_data = spark.read.format("image").load("D:\\data\\手写数字识别data\\test2")

# 2、取出图片的路径和数据，数据是二进制（字节数组）
path_data_df = image_data.select(col("image.origin").alias("path"), col("image.data").alias("data"))


# 3、编写自定义函数处理数据
def image_fun(data):
    # 将字节数组转换成普通列表
    data = [int(i) for i in data]

    # 将图片像素值列表转换成特征向量
    features = Vectors.dense(data)

    return features


# 注册自定义函数
image_udf = udf(image_fun, VectorUDT())

# 通过自定义函数处理数据，取出文件名称和图片数据
features_df = path_data_df.select("path", image_udf("data").alias("features"))

# 加载模型
model = LogisticRegressionModel.load("../../data/image_model")

# 预测分类
predict = model.transform(features_df)

predict.select("path", "prediction").show(n=100, truncate=False)
